#include "detector.h"

DetectNode::DetectNode(std::string engine_name, float conf_thresh, float nms_thresh) 
  : engine_name_(engine_name), kConfThresh(conf_thresh), kNmsThresh(nms_thresh)
{
  cudaSetDevice(kGpuId);
  deserialize_engine(engine_name_, &runtime, &engine, &context);
  
  CUDA_CHECK(cudaStreamCreate(&stream));
  // Init CUDA preprocessing
  cuda_preprocess_init(kMaxInputImageSize);
  prepare_buffers(engine, &gpu_buffers[0], &gpu_buffers[1], &cpu_output_buffer);

}

void DetectNode::prepare_buffers(ICudaEngine* engine, float** gpu_input_buffer, float** gpu_output_buffer, float** cpu_output_buffer) 
{
  assert(engine->getNbBindings() == 2);
  // In order to bind the buffers, we need to know the names of the input and output tensors.
  // Note that indices are guaranteed to be less than IEngine::getNbBindings()
  const int inputIndex = engine->getBindingIndex(kInputTensorName);
  const int outputIndex = engine->getBindingIndex(kOutputTensorName);
  assert(inputIndex == 0);
  assert(outputIndex == 1);
  // Create GPU buffers on device
  CUDA_CHECK(cudaMalloc((void**)gpu_input_buffer, kBatchSize * 3 * kInputH * kInputW * sizeof(float)));
  CUDA_CHECK(cudaMalloc((void**)gpu_output_buffer, kBatchSize * kOutputSize * sizeof(float)));

  *cpu_output_buffer = new float[kBatchSize * kOutputSize];
}

void DetectNode::infer(IExecutionContext& context, cudaStream_t& stream, void** gpu_buffers, float* output, int batchsize) 
{
  context.enqueue(batchsize, gpu_buffers, stream, nullptr);
  CUDA_CHECK(cudaMemcpyAsync(output, gpu_buffers[1], batchsize * kOutputSize * sizeof(float), cudaMemcpyDeviceToHost, stream));
  cudaStreamSynchronize(stream);
}

void DetectNode::deserialize_engine(std::string& engine_name, IRuntime** runtime, ICudaEngine** engine, IExecutionContext** context) 
{
  std::ifstream file(engine_name, std::ios::binary);
  if (!file.good()) {
    std::cerr << "--Read " << engine_name << " error!\n" << std::endl;
    assert(false);
  }
  size_t size = 0;
  file.seekg(0, file.end);
  size = file.tellg();
  file.seekg(0, file.beg);
  char* serialized_engine = new char[size];
  assert(serialized_engine);
  file.read(serialized_engine, size);
  file.close();

  *runtime = createInferRuntime(gLogger);
  assert(*runtime);
  *engine = (*runtime)->deserializeCudaEngine(serialized_engine, size);
  assert(*engine);
  *context = (*engine)->createExecutionContext();
  assert(*context);
  delete[] serialized_engine;
}

void DetectNode::inference(cv::Mat& img)
{
  cuda_preprocess(img.ptr(), img.cols, img.rows, gpu_buffers[0], kInputW, kInputH, stream);
    // Run inference
  auto start = std::chrono::system_clock::now();
  infer(*context, stream, (void**)gpu_buffers, cpu_output_buffer, kBatchSize);
  auto end = std::chrono::system_clock::now();
  std::cout << "inference time: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;

  // NMS
  std::vector<Detection> res;
  nms(res, cpu_output_buffer, kConfThresh, kNmsThresh);
  // Draw bounding boxes
  draw_bbox(img, res);
}

void DetectNode::inference(cv::Mat& img, std::vector<Detection>& res)
{
  cuda_preprocess(img.ptr(), img.cols, img.rows, gpu_buffers[0], kInputW, kInputH, stream);
    // Run inference
  
  auto start = std::chrono::high_resolution_clock::now();
  // auto start = std::chrono::system_clock::now();
  infer(*context, stream, (void**)gpu_buffers, cpu_output_buffer, kBatchSize);
  // auto end = std::chrono::system_clock::now();
  // cost_time_ = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
  auto end = std::chrono::high_resolution_clock::now();
  cost_time_ = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count(); 

  // NMS
  // std::vector<Detection> res;
  nms(res, cpu_output_buffer, kConfThresh, kNmsThresh);
  // // Draw bounding boxes
  // draw_bbox(img, res);
}

DetectNode::~DetectNode()
{
  // Release stream and buffers
  cudaStreamDestroy(stream);
  CUDA_CHECK(cudaFree(gpu_buffers[0]));
  CUDA_CHECK(cudaFree(gpu_buffers[1]));
  delete[] cpu_output_buffer;
  cuda_preprocess_destroy();
  // Destroy the engine
  context->destroy();
  engine->destroy();
  runtime->destroy();
}

